"""
PyG版本的图批处理操作，与dgl.batch函数兼容
"""

import torch
import torch_npu
from torch_npu.contrib import transfer_to_npu
from torch_geometric.data import Batch
from typing import List, Union
from ops.src.graph.graph import Graph


def batch(graphs: List[Union[Graph, 'torch_geometric.data.Data']]) -> Graph:
    """
    将多个图批处理为一个大图，与dgl.batch函数兼容
    
    Args:
        graphs: 要批处理的图列表
        
    Returns:
        批处理后的图
    """
    if not graphs:
        return Graph()
    
    # 自定义批处理逻辑
    batched_edge_index = []
    batched_node_data = {}
    batched_edge_data = {}
    
    node_offset = 0
    
    # 推断设备：从第一个图的数据中获取设备信息
    device = torch.device('cpu')  # 默认CPU
    if graphs and hasattr(graphs[0], 'ndata'):
        for key, value in graphs[0].ndata.items():
            if isinstance(value, torch.Tensor):
                device = value.device
                break
    
    for graph in graphs:
        # 处理边索引，加上节点偏移
        if hasattr(graph, 'edge_index'):
            edge_index = graph.edge_index + node_offset
            batched_edge_index.append(edge_index)
        
        # 处理节点数据
        if hasattr(graph, 'ndata'):
            for key, value in graph.ndata.items():
                if key not in batched_node_data:
                    batched_node_data[key] = []
                batched_node_data[key].append(value)
        elif hasattr(graph, '_node_data'):
            for key, value in graph._node_data.items():
                if key not in batched_node_data:
                    batched_node_data[key] = []
                batched_node_data[key].append(value)
        
        # 处理边数据
        if hasattr(graph, 'edata'):
            for key, value in graph.edata.items():
                if key not in batched_edge_data:
                    batched_edge_data[key] = []
                batched_edge_data[key].append(value)
        elif hasattr(graph, '_edge_data'):
            for key, value in graph._edge_data.items():
                if key not in batched_edge_data:
                    batched_edge_data[key] = []
                batched_edge_data[key].append(value)
        
        # 更新节点偏移
        if hasattr(graph, 'num_nodes') and callable(graph.num_nodes):
            node_offset += graph.num_nodes()
        elif hasattr(graph, '_num_nodes'):
            node_offset += graph._num_nodes
        elif hasattr(graph, 'number_of_nodes'):
            node_offset += graph.number_of_nodes()
        else:
            # 从边索引推断节点数
            if hasattr(graph, 'edge_index') and graph.edge_index.numel() > 0:
                node_offset += graph.edge_index.max().item() + 1
    
    # 创建批处理后的图
    if batched_edge_index:
        edge_index = torch.cat(batched_edge_index, dim=1)
        batched_graph = Graph((edge_index[0], edge_index[1]), num_nodes=node_offset)
    else:
        batched_graph = Graph(num_nodes=node_offset)
    
    # 设置批处理后的节点数据
    for key, values in batched_node_data.items():
        batched_graph.ndata[key] = torch.cat(values, dim=0)
    
    # 设置批处理后的边数据
    for key, values in batched_edge_data.items():
        batched_graph.edata[key] = torch.cat(values, dim=0)
    
    # 🔧 关键修复：创建正确的批处理信息，并确保设备一致
    # 为每个图的节点分配正确的批次索引
    batch_indices = []
    for graph_idx, graph in enumerate(graphs):
        if hasattr(graph, 'num_nodes') and callable(graph.num_nodes):
            num_nodes = graph.num_nodes()
        elif hasattr(graph, '_num_nodes'):
            num_nodes = graph._num_nodes
        elif hasattr(graph, 'number_of_nodes'):
            num_nodes = graph.number_of_nodes()
        else:
            # 从节点数据推断
            if batched_node_data and 'attr' in batched_node_data:
                num_nodes = batched_node_data['attr'][graph_idx].shape[0]
            else:
                num_nodes = 1  # 默认值
        
        # 为这个图的所有节点分配相同的批次索引，确保设备一致
        batch_indices.append(torch.full((num_nodes,), graph_idx, dtype=torch.long, device=device))
    
    if batch_indices:
        batched_graph.batch = torch.cat(batch_indices, dim=0)
        pass  # batch信息已设置
    
    return batched_graph 